package com.jsean.demo;

import com.aayushatharva.brotli4j.Brotli4jLoader;
import com.aayushatharva.brotli4j.decoder.BrotliInputStream;
import com.aayushatharva.brotli4j.encoder.BrotliOutputStream;
import com.aayushatharva.brotli4j.encoder.Encoder;
import com.github.luben.zstd.Zstd;
import net.jpountz.lz4.LZ4Compressor;
import net.jpountz.lz4.LZ4Factory;
import net.jpountz.lz4.LZ4FastDecompressor;
import org.openjdk.jmh.annotations.*;
import org.openjdk.jmh.runner.Runner;
import org.openjdk.jmh.runner.RunnerException;
import org.openjdk.jmh.runner.options.Options;
import org.openjdk.jmh.runner.options.OptionsBuilder;
import org.xerial.snappy.Snappy;

import java.io.*;
import java.util.concurrent.TimeUnit;
import java.util.zip.GZIPInputStream;
import java.util.zip.GZIPOutputStream;

@State(Scope.Thread)
@BenchmarkMode(Mode.AverageTime)
@OutputTimeUnit(TimeUnit.MILLISECONDS)
@Warmup(iterations = 3, time = 1)
@Measurement(iterations = 5, time = 1)
@Fork(1)
public class BenchmarkPlatform {

    @Param({"lorem.txt", "json_sample.txt", "binary_data.bin"})
    private String testFile;

    private byte[] originalData;
    private byte[] lz4Compressed;
    private byte[] zstdCompressed;
    private byte[] snappyCompressed;
    private byte[] gzipCompressed;
    private byte[] brotliCompressed;

    private LZ4Compressor lz4Compressor;
    private LZ4FastDecompressor lz4Decompressor;

    @Setup(Level.Trial)
    public void setup() throws IOException {
        // 加载测试数据
        originalData = loadTestData(testFile);

        // 初始化压缩器
        LZ4Factory factory = LZ4Factory.fastestInstance();
        lz4Compressor = factory.fastCompressor();
        lz4Decompressor = factory.fastDecompressor();

        // 预加载 Brotli
        Brotli4jLoader.ensureAvailability();

        // 预先压缩数据
        lz4Compressed = compressLZ4(originalData);
        zstdCompressed = compressZstd(originalData);
        snappyCompressed = compressSnappy(originalData);
        gzipCompressed = compressGzip(originalData);
        brotliCompressed = compressBrotli(originalData);
    }

    private byte[] loadTestData(String filename) throws IOException {
        try (InputStream is = getClass().getClassLoader().getResourceAsStream(filename);
             ByteArrayOutputStream bos = new ByteArrayOutputStream()) {
            byte[] buffer = new byte[1024];
            int len;
            while ((len = is.read(buffer)) > -1) {
                bos.write(buffer, 0, len);
            }
            return bos.toByteArray();
        }
    }

    // ========== LZ4 基准测试 ==========
    @Benchmark
    public byte[] lz4Compress() {
        return compressLZ4(originalData);
    }

    private byte[] compressLZ4(byte[] data) {
        int maxCompressedLength = lz4Compressor.maxCompressedLength(data.length);
        byte[] compressed = new byte[maxCompressedLength];
        int compressedLength = lz4Compressor.compress(data, 0, data.length, compressed, 0, maxCompressedLength);
        byte[] result = new byte[compressedLength];
        System.arraycopy(compressed, 0, result, 0, compressedLength);
        return result;
    }

    @Benchmark
    public byte[] lz4Decompress() {
        return decompressLZ4(lz4Compressed, originalData.length);
    }

    private byte[] decompressLZ4(byte[] compressed, int originalLength) {
        byte[] restored = new byte[originalLength];
        lz4Decompressor.decompress(compressed, 0, restored, 0, originalLength);
        return restored;
    }

    // ========== Zstandard 基准测试 ==========
    @Benchmark
    public byte[] zstdCompress() {
        return compressZstd(originalData);
    }

    private byte[] compressZstd(byte[] data) {
        return Zstd.compress(data);
    }

    @Benchmark
    public byte[] zstdDecompress() {
        return decompressZstd(zstdCompressed, originalData.length);
    }

    private byte[] decompressZstd(byte[] compressed, int originalLength) {
        try {
            return Zstd.decompress(compressed, originalLength);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    // ========== Snappy 基准测试 ==========
    @Benchmark
    public byte[] snappyCompress() {
        return compressSnappy(originalData);
    }

    private byte[] compressSnappy(byte[] data) {
        try {
            return Snappy.compress(data);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Benchmark
    public byte[] snappyDecompress() {
        return decompressSnappy(snappyCompressed);
    }

    private byte[] decompressSnappy(byte[] compressed) {
        try {
            return Snappy.uncompress(compressed);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    // ========== GZIP 基准测试 ==========
    @Benchmark
    public byte[] gzipCompress() {
        return compressGzip(originalData);
    }

    private byte[] compressGzip(byte[] data) {
        try (ByteArrayOutputStream bos = new ByteArrayOutputStream();
             GZIPOutputStream gzip = new GZIPOutputStream(bos)) {
            gzip.write(data);
            gzip.finish();
            return bos.toByteArray();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Benchmark
    public byte[] gzipDecompress() {
        return decompressGzip(gzipCompressed);
    }

    private byte[] decompressGzip(byte[] compressed) {
        try (ByteArrayInputStream bis = new ByteArrayInputStream(compressed);
             GZIPInputStream gzip = new GZIPInputStream(bis);
             ByteArrayOutputStream bos = new ByteArrayOutputStream()) {
            byte[] buffer = new byte[1024];
            int len;
            while ((len = gzip.read(buffer)) > 0) {
                bos.write(buffer, 0, len);
            }
            return bos.toByteArray();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    // ========== Brotli 基准测试 ==========
    @Benchmark
    public byte[] brotliCompress() {
        return compressBrotli(originalData);
    }

    private byte[] compressBrotli(byte[] data) {
        try (ByteArrayOutputStream bos = new ByteArrayOutputStream();
             BrotliOutputStream brotli = new BrotliOutputStream(bos, new Encoder.Parameters().setQuality(4))) {
            brotli.write(data);
            brotli.flush();
            return bos.toByteArray();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Benchmark
    public byte[] brotliDecompress() {
        return decompressBrotli(brotliCompressed);
    }

    private byte[] decompressBrotli(byte[] compressed) {
        try (ByteArrayInputStream bis = new ByteArrayInputStream(compressed);
             BrotliInputStream brotli = new BrotliInputStream(bis);
             ByteArrayOutputStream bos = new ByteArrayOutputStream()) {
            byte[] buffer = new byte[1024];
            int len;
            while ((len = brotli.read(buffer)) > 0) {
                bos.write(buffer, 0, len);
            }
            return bos.toByteArray();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    public static void main(String[] args) throws RunnerException {
        Options opt = new OptionsBuilder()
                .include(BenchmarkPlatform.class.getSimpleName())
                .build();

        new Runner(opt).run();
    }
}